import os
import json
from PIL import Image
from tqdm import tqdm

from config import dataset_dir, figure_save_dir, label_save_path
from utils.bbox import IOU, yolo2bbox 

YOLO_FIGURE_TYPE = 1
YOLO_SUBFIGURE_TYPE = 2

'''
所有存储的bbox格式为 l ,t ,r, b
'''

class Figure:
    def __init__(self, figure_bbox, all_subfigure_bbox_list):
        self.figure_bbox = figure_bbox
        self.subfigure_bbox_list = []
        self.checkoutSubfigure(all_subfigure_bbox_list)
    
    def checkoutSubfigure(self, all_subfigure_bbox_list):
        for subfigure_bbox in all_subfigure_bbox_list:
            iou = IOU(self.figure_bbox, subfigure_bbox)
            if iou != 0:
                self.subfigure_bbox_list.append(subfigure_bbox)        
                    
    def hasSubfigure(self):
        return len(self.subfigure_bbox_list) > 0
    
    def format(self):
        return [self.figure_bbox, self.subfigure_bbox_list]


class Preprocess:
    def __init__(self, dataset_dir, figure_save_dir, label_save_path):
        self.dataset_dir = dataset_dir
        self.figure_save_dir = figure_save_dir
        self.label_save_path = label_save_path
        self.has_subfigure_figure_list = dict()

    def run(self):
        for item in tqdm(os.listdir(self.dataset_dir)):
            # if item != "3":
            #     continue
            obj_dir = os.path.join(self.dataset_dir, item)
            for filename in os.listdir(obj_dir):
                # if filename != "3_000044.jpg":
                #     continue
                filepath = os.path.join(obj_dir, filename)
                if os.path.isfile(filepath):
                    labelpath = self.getLabelPath(obj_dir, filename)
                    if os.path.isfile(labelpath):
                        self.dealOneImage(filepath, labelpath)
        self.save()
    
    def dealOneImage(self, filepath, labelpath):
        img = Image.open(filepath)
        w, h = img.size
        figures, subfigures = self.dealLabel(labelpath, w, h)
        idx = 0
        image_name = filepath.split("/")[-1]
        prefix = image_name.split(".")[0]
        
        for figure in figures:
            figure_object = Figure(figure, subfigures)  
            if figure_object.hasSubfigure():
                # 确定裁剪图片的名字
                cut_name = prefix + "_" +str(idx) + ".jpg"
                idx += 1
                # 需要裁剪图片并保存
                cut_image = img.crop(figure)
                cut_path = os.path.join(self.figure_save_dir, cut_name)
                print(figure)
                print(cut_name)
                cut_image.save(cut_path)
                # 保存标注
                self.has_subfigure_figure_list[cut_name] = figure_object.format()
                
    @staticmethod
    def getLabelPath(root, filename):
        prefix = filename.split(".")[0]
        labelname = prefix + ".txt"
        labelpath = os.path.join(root, "classes_labels")
        labelpath = os.path.join(labelpath, labelname)
        return labelpath

    @staticmethod
    def dealLabel(labelpath, img_w, img_h):
        figure_bbox_list = []
        subfigure_bbox_list = []
        with open(labelpath, "r") as fp:
            while True:
                line = fp.readline()
                line = line.strip()
                if len(line) == 0:
                    break
                data = line.split()
                typ = int(data[0])
                if typ != YOLO_FIGURE_TYPE and typ != YOLO_SUBFIGURE_TYPE:
                    continue
                cx = int(float(data[1]) * img_w)
                cy = int(float(data[2]) * img_h)
                w = int(float(data[3]) * img_w)
                h = int(float(data[4]) * img_h)
                l, t, r, b = yolo2bbox(cx, cy, w, h)
                bbox = [l, t, r, b]
                if typ == YOLO_FIGURE_TYPE:
                    figure_bbox_list.append(bbox)
                else:
                    subfigure_bbox_list.append(bbox)
        return figure_bbox_list, subfigure_bbox_list
    
    def save(self):
        with open(self.label_save_path, "w") as fp:
            json.dump(self.has_subfigure_figure_list, fp)

if __name__ == "__main__":
    preprocess = Preprocess(dataset_dir, figure_save_dir, label_save_path)
    preprocess.run()
    


